package edu.cornell.cs4740.sentencegenerator;

import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;

public class UnigramTable {
  private Map<String, Long> unigramTable;
  private Map<String, Double> unsmoothedProbTable;
  private Map<Integer, Double> gtProbTable;
  private long numberOfUnigrams = 0;
  private double vocabSize = 0;
  private int maxFreq = 0;

  public UnigramTable(List<String> words, boolean unknownWords) {
    unigramTable = new HashMap<String, Long>();
    HashSet<String> wordsSeen = new HashSet<String>();

    for (String w : words) {
      String word = w;
      if (unknownWords) {
        if (!wordsSeen.contains(word)) {
          wordsSeen.add(word);
          word = Utilities.UNKNOWN;
        }
      }

      if (unigramTable.containsKey(word)) {
        long wordCount = unigramTable.get(word);
        wordCount++;
        unigramTable.put(word, new Long(wordCount));
        maxFreq = (wordCount > maxFreq) ? (int) wordCount : maxFreq;
      } else {
        unigramTable.put(word, new Long(1));
        maxFreq = (1 > maxFreq) ? 1 : maxFreq;
      }

      numberOfUnigrams++;
    }
    vocabSize = unigramTable.size();
  }

  public Map<String, Double> generateProbabilityTable() {
    Map<String, Double> probabilityTable = new HashMap<String, Double>();

    for (String word : unigramTable.keySet()) {
      long wordCount = unigramTable.get(word);
      double wordProbability = ((double) wordCount)
          / ((double) numberOfUnigrams);
      probabilityTable.put(word, new Double(wordProbability));
    }

    unsmoothedProbTable = probabilityTable;
    return probabilityTable;
  }

  public double getUnsmoothedProb(String word) {
    if (unsmoothedProbTable == null) {
      generateProbabilityTable();
    }

    Double ans = unsmoothedProbTable.get(word);
    if (ans == null) {
      ans = unsmoothedProbTable.get(Utilities.UNKNOWN);
      if (ans == null) {
        return 0;
      }
    }
    return ans;
  }

  public double getLaplacianProb(String word) {
    Long count = unigramTable.get(word);
    if (count == null) {
      count = unigramTable.get(Utilities.UNKNOWN);
      if (count == null) {
        count = new Long(0);
      }
    }

    return (count + 1) / (numberOfUnigrams + vocabSize);
  }

  private void makeGoodTuringTable() {
    int[] freqOfFreq = new int[maxFreq + 1];
    for (int i = 0; i <= maxFreq; i++) {
      freqOfFreq[i] = 0;
    }

    for (String key : unigramTable.keySet()) {
      Long tempCount = unigramTable.get(key);
      freqOfFreq[tempCount.intValue()] += 1;
    }

    gtProbTable = new HashMap<Integer, Double>();

    for (int k = 0; k <= maxFreq; k++) {
      if (freqOfFreq[k] == 0 && k != 0) {
        continue;
      }

      int next = (k == maxFreq) ? maxFreq : k + 1;
      while (k < maxFreq) {
        if (freqOfFreq[next] != 0) {
          break;
        }
        next++;
      }

      double newProb;
      if (k == 0) {
        newProb = (k + 1) * (double) freqOfFreq[next] / (double) freqOfFreq[1];
      } else if (k <= 5) {
        newProb = (k + 1) * (double) freqOfFreq[next] / (double) freqOfFreq[k];
      } else {
        newProb = (double) k;
      }

      gtProbTable.put(k, newProb / unigramTable.size());
    }

//    double sumXes = 0; // sum(x_1,...,x_n)
//    double sumYs = 0; // sum(y_1,...,y_n)
//    double sumXYs = 0; // sum(xy_1,...,xy_n)
//    double sumSquareXes = 0; // sum(x^2)
//
//    for (int xKey = 1; xKey < freqOfFreq.length; xKey++) {
//      double x = Math.log(xKey);
//      double y = freqOfFreq[xKey];
//      sumXes += x;
//      sumYs += y;
//      sumXYs += x * y;
//      sumSquareXes += x * x;
//    }
//
//    double n = freqOfFreq.length;
//
//    double beta = sumXYs - ((sumXes * sumYs) / n);
//    beta = beta / (sumSquareXes - ((sumXes * sumXes) / n));
//
//    double alpha = (sumYs - (beta * sumXes)) / n;
//
//    gtProbTable = new HashMap<Integer, Double>();
//    for (int k = 0; k <= maxFreq; k++) {
//      double log = (k == 0) ? Math.log(1) : Math.log(k);
//      double nc = alpha + (beta * log);
//      nc = Math.pow(Math.E, nc);
//      double nc1 = alpha + (beta * Math.log(k + 1));
//      nc1 = Math.pow(Math.E, nc1);
//      double newProb = (k + 1) * nc1;
//      newProb = newProb / (nc);
//      gtProbTable.put(k, newProb / unigramTable.size());
//    }
  }

  public double getGoodTuringProb(String word) {
    if (gtProbTable == null) {
      makeGoodTuringTable();
    }

    String w1 = (unigramTable.containsKey(word)) ? word : Utilities.UNKNOWN;

    Long count = unigramTable.get(word);
    if (count == null) {
      return gtProbTable.get(0);
    }

    return gtProbTable.get(count.intValue());
  }
}
